#include <strmif.h>
#include "mkvcontext.hpp"
#include "mkvstreamvideompeg2.hpp"
#include <climits>
#include <uuids.h>
#include <dvdmedia.h>
//#ifdef _DEBUG
//#include "odbgstream.hpp"
//using std::endl;
//#endif

namespace MkvMux
{


MkvStreamVideoMpeg2::Mpeg2VideoFrame::Mpeg2VideoFrame(MPEGFrame* f) :
   m_pFrame(f)
{
}


MkvStreamVideoMpeg2::Mpeg2VideoFrame::~Mpeg2VideoFrame()
{
   delete m_pFrame;
}


MkvStreamVideoMpeg2::Type
MkvStreamVideoMpeg2::Mpeg2VideoFrame::type() const
{
   switch (m_pFrame->frameType)
   {
      case 'I':
      default:
         return kTypeIFrame;
         
      case 'P':
         return kTypePFrame;
         
      case 'B':
         return kTypeBFrame;
   }
}


bool MkvStreamVideoMpeg2::Mpeg2VideoFrame::IsKey() const
{
    return (type() == kTypeIFrame);
}

         

ULONG MkvStreamVideoMpeg2::Mpeg2VideoFrame::size() const
{
   return m_pFrame->size;
}

   
const BYTE* MkvStreamVideoMpeg2::Mpeg2VideoFrame::data() const
{
   return m_pFrame->data;
}


ULONG MkvStreamVideoMpeg2::Mpeg2VideoFrame::curr_timecode_ms() const
{
   return m_pFrame->timecode_ms;
}


SHORT MkvStreamVideoMpeg2::Mpeg2VideoFrame::prev_timecode_ms() const
{
   const int result = m_pFrame->firstRef_ms - m_pFrame->timecode_ms;
   assert(result <= 0);
   assert(result >= SHRT_MIN);
   
   return static_cast<SHORT>(result);
}


SHORT MkvStreamVideoMpeg2::Mpeg2VideoFrame::next_timecode_ms() const
{
   const int result = m_pFrame->secondRef_ms - m_pFrame->timecode_ms;
   assert(result >= 0);
   assert(result <= SHRT_MAX);
   
   return static_cast<SHORT>(result);
}


MkvStreamVideoMpeg2::MkvStreamVideoMpeg2(
    MkvContext& c,
    const AM_MEDIA_TYPE& mt) :
    MkvStreamVideo(c, mt)
{
}


void MkvStreamVideoMpeg2::GetSequenceHeader(
    const BYTE*& h,
    ULONG& cb) const
{
   const AM_MEDIA_TYPE& mt = m_mt;
   assert(mt.majortype == MEDIATYPE_Video);
   assert(mt.subtype == MEDIASUBTYPE_MPEG2_VIDEO);  //TODO: more subtypes here
   assert(mt.cbFormat >= sizeof(VIDEOINFOHEADER2));
   assert(mt.pbFormat);  //fully vetted in InpinVideo::QueryAccept

   const MPEG2VIDEOINFO& mpg = (MPEG2VIDEOINFO&)(*mt.pbFormat);
   assert(mt.cbFormat >= SIZE_MPEG2VIDEOINFO(&mpg));

   cb = mpg.cbSequenceHeader;
   assert(cb >= 4);

   const VIDEOINFOHEADER2& vih = mpg.hdr;
   const BITMAPINFOHEADER& bmih = vih.bmiHeader;
   bmih;

   //http://dvd.sourceforge.net/dvdinfo/mpeghdrs.html#seq

   h = MPEG2_SEQUENCE_INFO(&mpg);
   assert(h[0] == 0);
   assert(h[1] == 0);
   assert(h[2] == 1);
   assert(h[3] == 0xB3);
}



void MkvStreamVideoMpeg2::WriteTrackName()
{
    EbmlIO::File& f = m_context.m_file;

    f.WriteID2(0x536E);   //name
    f.Write1UTF8(L"MPEG-2 video stream");
}



void MkvStreamVideoMpeg2::WriteTrackCodecID()
{
    EbmlIO::File& f = m_context.m_file;

    f.WriteID1(0x86);  //Codec ID
    f.Write1String("V_MPEG2");
}


void MkvStreamVideoMpeg2::WriteTrackCodecPrivate()
{
    EbmlIO::File& f = m_context.m_file;
    
    const BYTE* h;
    ULONG cb;

    GetSequenceHeader(h, cb);
    assert(h);
    assert(cb >= 4);
    
    f.WriteID2(0x63A2);  //Codec Private
    f.Write4UInt(cb);
    f.Write(h, cb);
}
    

void MkvStreamVideoMpeg2::WriteTrackCodecName()
{
    EbmlIO::File& f = m_context.m_file;

    f.WriteID3(0x258688);  //Codec Name
    f.Write1UTF8(L"MPEG-2");
}


void MkvStreamVideoMpeg2::WriteTrackSettings()
{
    EbmlIO::File& f = m_context.m_file;

    f.WriteID1(0xE0);  //video settings
    
    //allocate 2 bytes of storage for size of settings
    const __int64 begin_pos = f.SetPosition(2, STREAM_SEEK_CUR);
    
    const BYTE* ha;
    ULONG hn;
    
    GetSequenceHeader(ha, hn);
    assert(ha);
    assert(hn >= 4);
    
    binary* const data = new binary[hn];
    memcpy(data, ha, hn);
    
    MPEGChunk chunk(data, hn);  //deallocates data in dtor
    
    MPEG2SequenceHeader h;    
    ParseSequenceHeader(&chunk, h);

    const BITMAPINFOHEADER& bmih = GetBitmapInfoHeader();
    assert(bmih.biSize >= sizeof(BITMAPINFOHEADER));
    assert(bmih.biWidth > 0);
    assert(bmih.biHeight > 0);
    assert(h.width == ULONG(bmih.biWidth));
    assert(h.height == ULONG(bmih.biHeight));

    assert(h.width <= USHRT_MAX);
    const USHORT width = static_cast<USHORT>(h.width);
    
    assert(h.height <= USHRT_MAX);
    const USHORT height = static_cast<USHORT>(h.height);

    f.WriteID1(0xB0);  //width
    f.Write1UInt(2);
    f.Serialize2UInt(width);
    
    f.WriteID1(0xBA);  //height
    f.Write1UInt(2);
    f.Serialize2UInt(height);
    
    //TODO: the BMIH tells us whether the stride is different from width;
    //is that the pixel crop values here?

    //TODO:    
    //mkvmerge sets to Interlaced flag for the video track to 0.
    //Is this value determined from h.progressiveSequence?
    
    //aspect ratio = 4:3
    //720:480 = 1.5 = 3:2 = 4.5:3
    //720 x factor / 480 = 4:3
    //factor = (4/3) * (480/720) = 
    
    f.WriteID3(0x2383E3);  //frame rate
    f.Write1UInt(4);
    f.Serialize4Float(h.frameRate);    

    const __int64 end_pos = f.GetPosition();
    
    const __int64 size_ = end_pos - begin_pos;
    assert(size_ <= USHRT_MAX);
    
    const USHORT size = static_cast<USHORT>(size_);

    f.SetPosition(begin_pos - 2);
    f.Write2UInt(size);
    
    f.SetPosition(end_pos);
}


HRESULT MkvStreamVideoMpeg2::Receive(IMediaSample* pSample)
{
    assert(pSample);

    if (m_context.m_file.GetStream() == 0)
        return S_OK;  //TODO: pass EOS downstream immediately

    BYTE* ptr;

    const HRESULT hr = pSample->GetPointer(&ptr);
    assert(SUCCEEDED(hr));
    assert(ptr);
    hr;
         
    const long len = pSample->GetActualDataLength();
    assert(len >= 0);

    m_parser.WriteData(ptr, len);

    while (MPEGFrame* f = m_parser.ReadFrame())
        CreateNewFrame(f);
     
    return S_OK;
}


int MkvStreamVideoMpeg2::EndOfStream()
{
    if (m_context.m_file.GetStream())
    {
        m_parser.SetEOS();

        while (MPEGFrame* f = m_parser.ReadFrame())
            CreateNewFrame(f);
    }
         
    return m_context.NotifyVideoEOS(this);
}


#if 0 //def _DEBUG
static std::ostream& operator<<(std::ostream& os, const MPEGFrame* f)
{
    assert(f);
    
    os << "(" << f->frameType << ") tc[ms]=" << f->timecode_ms;
    
    if (f->frameType != 'I')
        os << " back[ms]=" << f->firstRef_ms;
        
    if (f->frameType == 'B')
        os << " fwd[ms]=" << f->secondRef_ms;
        
    //os << " dur[ms]=" << f->duration_ms;
            
    return os;
}
#endif


void MkvStreamVideoMpeg2::CreateNewFrame(MPEGFrame* f)
{
   assert(f);
   assert(f->timecode_ms >= 0);

//#ifdef _DEBUG
//   odbgstream os;
//   os << "Mpeg2::WriteFrame: " << f << endl;
//#endif

    Mpeg2VideoFrame* const pFrame = new (std::nothrow) Mpeg2VideoFrame(f);
    assert(pFrame);  //TODO

    assert(!m_vframes.empty() || (pFrame->type() != kTypeBFrame));
           
    m_vframes.push_back(pFrame);

    switch (pFrame->type())
    {
        case kTypeIFrame:
        case kTypePFrame:
            m_context.NotifyVideoFrame(this, pFrame);
            break;
            
        case kTypeBFrame:
        default:
            break;            
    }
}


void MkvStreamVideoMpeg2::Mpeg2VideoFrame::Write(
    const MkvStream& s,
    ULONG cluster_timecode_ms) const
{
    EbmlIO::File& file = s.m_context.m_file;
    
    //block = 1 byte ID + 4 byte size + f->size
    //block duration = 1 byte ID + 1 byte size + 1 byte value
    //reference block = 1 byte ID + 2 byte size + 2(?) byte signed value

    const ULONG block_size = 1 + 2 + 1 + size();      //tn, tc, flg, f
    ULONG block_group_size = 1 + 4 + block_size; //block id, size, payload

    if (type() != kTypeIFrame)
        block_group_size += 4;  //for back ref
     
    if (type() == kTypeBFrame)
        block_group_size += 4;  //for fwd ref

    file.WriteID1(0xA0);                //block group ID
    file.Write4UInt(block_group_size);  //size of payload for this block group    

#ifdef _DEBUG    
    const __int64 pos = file.GetPosition();
#endif

    //begin block

    //It appears that the WD TV HD Media Player (Gen 1) does not 
    //like SimpleBlock elements.  In fact it corrupts the device
    //so badly you must unplug the unit to restore its ability
    //to render MKV files.  So for now we'll continue to use 
    //standard Block elements.  (If you're on the engineering team
    //for the WD TV HD Media Player, please contact me about 
    //your preferred format of MKV files.)
    //
    //file.WriteID1(0xA3);    //SimpleBlock ID
    file.WriteID1(0xA1);  //(non-Simple) Block ID
    file.Write4UInt(block_size);

    const int tn_ = s.GetTrackNumber();
    assert(tn_ > 0);
    assert(tn_ <= 127);

    const BYTE tn = static_cast<BYTE>(tn_);

    file.Write1UInt(tn);   //track number

    const ULONG ft = curr_timecode_ms();
      
    {
        const LONG tc_ = LONG(ft) - LONG(cluster_timecode_ms);
        assert(tc_ >= SHRT_MIN);
        assert(tc_ <= SHRT_MAX);

        const SHORT tc = static_cast<SHORT>(tc_);

        file.Serialize2SInt(tc);       //relative timecode 
    }
   
    const BYTE flags = 0;

    //See comments above about WD TV HD Media Player.    
    //if (type() == kTypeIFrame)
    //    flags |= BYTE(1 << 7);    //SimpleBlock only?
        
    file.Write(&flags, 1);   //written as binary, not uint

    file.Write(data(), size());  //frame

   //end block

    if (type() != kTypeIFrame)
    {
        file.WriteID1(0xFB);
        file.Write1UInt(2);
        file.Serialize2SInt(prev_timecode_ms());

        if (type() == kTypeBFrame)
        {
             file.WriteID1(0xFB);
             file.Write1UInt(2);
             file.Serialize2SInt(next_timecode_ms());
        } 
    }

#ifdef _DEBUG
    const __int64 newpos = file.GetPosition();
    assert((newpos - pos) == block_group_size);
#endif

    
}


LONG MkvStreamVideoMpeg2::last_timecode_ms() const
{
    if (m_vframes.empty())
        return -1;
        
    typedef MkvStreamVideo::frames_t::const_iterator iter_t;

    const iter_t i = m_vframes.begin();
    iter_t j = m_vframes.end();
        
    for (;;)    
    {
        const VideoFrame* const pvf_ = *--j;
        assert(pvf_);
        
        const Mpeg2VideoFrame* const pvf =
            static_cast<const Mpeg2VideoFrame*>(pvf_);
        
        if (pvf->type() != kTypeBFrame)
            return pvf->curr_timecode_ms();
            
        if (i == j)  //should never happen
            return -1;
    }
}


void MkvStreamVideoMpeg2::Flush()
{
    if (m_context.m_file.GetStream())
    {
        m_parser.SetEOS();

        while (MPEGFrame* f = m_parser.ReadFrame())
            CreateNewFrame(f);
    }

    MkvStreamVideo::Flush();
}

}  //end namespace MkvMux

